import datetime
from app.db.models.asset_tasks import AssetTasks
from app.db.models.jobs import Job
from app.db.models.wass import ScanEvents
from app.libs.score import Score


class JobSecurityStatus:

    def caculate_score(self, all_warnings, update_time):
        update_dict = {'score': 0, "level": "未知", "update_time": ""}
        if all_warnings:
            score_dict = Score.score_info(event_map={}, warnings=all_warnings)
            update_dict['score'] = score_dict.get('score', 0)
            update_dict['level'] = score_dict.get('level', "未知")

        if not all_warnings and all_warnings != '':
            update_dict['level'] = '安全'

        if update_time:
            update_dict['update_time'] = update_time

        return update_dict

    def web_security_status(self):
        task_types = ['vul', 'ssl', 'securityEvent', 'content']
        lookups = [{'$lookup': {'from': "tasks", 'localField': f"{taskType}ResultId", 'foreignField': "_id",
                                'as': f"{taskType}"}} for taskType in task_types]

        _projects = ['{}.endTime', '{}.result']
        project = {_project.format(task_type): 1 for _project in _projects for task_type in task_types}
        aggregate = [
            {'$match': {"verification.isVerified": True, "assetType": {'$ne': 'host'}}},
            *lookups,
            {'$project': project},
        ]
        jobs = Job.objects.aggregate(aggregate)
        for job in jobs:
            update_time = ''
            all_warnings = []
            is_checked = False
            for task_type in task_types:
                task_type_data = job.get(task_type, [])
                for item in task_type_data:
                    is_checked = True
                    end_time = item.get('endTime')
                    if (not update_time and end_time) or (end_time and end_time > update_time):
                        update_time = end_time

                    warnings = item.get('result', {}).get('warnings', [])
                    all_warnings.extend(warnings)

            if not is_checked:
                continue

            update_dict = self.caculate_score(all_warnings, update_time)
            Job.objects(pk=job['_id']).update_one(securityStatus=update_dict)

    def host_security_status(self):
        query_dict = {"verification.isVerified": True, "assetType": "host"}
        jobs = Job.objects.find(query_dict)

        total = 0
        job_id_list = []
        job_warning_maps = dict()
        for job in jobs:
            job_id_list.append(job.id)
            job_warning_maps[job.id] = {"update_time": "", "warnings": [], "has_task": False}

        aggregate = [
            {'$match': {'taskType': "hostVul", "jobId": {"$in": job_id_list}, "hostVulResultId": {"$exists": True}}},
            {'$group': {'_id': '$jobId', 'id': {'$last': '$_id'}, 'hostVulResultId': {'$last': '$hostVulResultId'}}},
            {'$lookup': {'from': "tasks", 'localField': "hostVulResultId", 'foreignField': "_id", 'as': "task"}},
        ]

        asset_tasks = AssetTasks.objects.aggregate(aggregate)
        for asset_task in asset_tasks:
            if not asset_task.get('task'):
                continue

            task = asset_task.get('task')[0]
            job_id = task.get('jobId')
            job_warning_maps[job_id]['has_task'] = True
            job_warning_maps[job_id]["update_time"] = task.get("endTime", "")

            query = {"task_session_id": task.get("taskSessionId"), "event_name": "network_vuln"}
            host_vul_events = ScanEvents.objects.find(query)
            for event in host_vul_events:
                warning = {"level": event.severity}
                job_warning_maps[job_id]["warnings"].append(warning)

        for job_id, data in job_warning_maps.items():
            has_task = data.get("has_task", [])
            if not has_task:
                continue

            warnings = data.get("warnings", [])
            update_time = data.get("update_time", "")

            update_dict = self.caculate_score(warnings, update_time)
            Job.objects(pk=job_id).update_one(securityStatus=update_dict)

        return total, job_warning_maps

    def service(self):
        print(f"{datetime.datetime.now()} 开始同步......")
        self.web_security_status()
        self.host_security_status()
        print(f"{datetime.datetime.now()} 同步结束......")
